#include <tinx/debug.h>
#include <tinx/interrupt.h>
#include <tinx/global.h>
#include <tinx/task.h>
#include <tinx/memory.h>
#include <tinx/printk.h>
#include <tinx/fs.h>
#include <tinx/user.h>
#include <tinx/cpu.h>
#include <tinx/string.h>
#include <tinx/arena.h>
#include <tinx/list.h>
#include <tinx/assert.h>
#include <tinx/types.h>
#include <tinx/syscall.h>
#include <tinx/stdlib.h>

task_t *idle_task;
task_t *task_table[TASK_NR];

static list_t schedule_list;

task_t *running_task()
{
    asm volatile(
        "movl %esp, %eax\n"
        "andl $0xfffff000, %eax");
}

pid_t task_getpid()
{
    task_t *task = running_task();
    return task->pid;
}

pid_t task_getppid()
{
    task_t *task = running_task();
    return task->ppid;
}

pid_t task_getuid()
{
    task_t *task = running_task();
    return task->uid;
}

task_t *get_free_task()
{
    task_t *task = NULL;
    for (int i = 1; i < TASK_NR; i++)
    {
        if (task_table[i])
            continue;
        task = (task_t *)alloc_kpage(1);
        task_table[i] = task;
        task->pid = i;
        break;
    }
    return task;
}

pid_t task_remove(task_t *task)
{
    assert(task);

    pid_t pid = task->pid;
    assert(task_table[pid] == task);

    free_kpage((u32)task, 1);
    task_table[pid] = NULL;
    return pid;
}

extern void interrupt_exit();

void task_to_user_mode(handler_t target)
{
    task_t *task = running_task();

    task->vmap = kmalloc(sizeof(bitmap_t));
    bitmap_init(task->vmap, (char *)alloc_kpage(1), PAGE_SIZE, IDX(USER_MMAP_ADDR));

    task->pgdir = (u32)copy_pgdir();
    task_activity(task);

    u32 addr = (u32)task + PAGE_SIZE;
    addr -= sizeof(intr_frame_t);

    intr_frame_t *iframe = (intr_frame_t *)addr;
    iframe->eip = (u32)target;
    iframe->cs = USER_CODE_SELECTOR;
    iframe->ds = USER_DATA_SELECTOR;
    iframe->es = USER_DATA_SELECTOR;
    iframe->fs = USER_DATA_SELECTOR;
    iframe->gs = USER_DATA_SELECTOR;
    iframe->ss = USER_DATA_SELECTOR;

    iframe->esp = USER_MEMORY_END;
    iframe->error = 0x20250601;
    iframe->eflags = (0 << 12 | 0b10 | 1 << 9);

    asm volatile(
        "movl %0, %%esp\n"
        "jmp interrupt_exit\n" :: "m"(iframe) : "memory");
}

extern file_t file_table[];

task_t *task_create(char *name, u8 prio, handler_t handler, uid_t uid)
{
    task_t *task = get_free_task();
    
    strncpy(task->name, name, 16);
    task->state = TASK_READY;
    task->jiffies = 0;
    task->prio = prio;
    task->ticks = prio;
    task->flags = 0;
    task->stack = (u32)task + PAGE_SIZE;
    task->tls = NULL;

    task->ppid = 0;
    task->uid = 0;
    task->exittid = NULL;

    task->stack -= sizeof(intr_frame_t);
    task->stack -= sizeof(task_frame_t);

    task_frame_t *frame = (task_frame_t *)task->stack;
    frame->eip = handler;
    frame->esi = 0x11111111;
    frame->edi = 0x22222222;
    frame->ebx = 0x33333333;
    frame->ebp = 0x44444444;

    task->pgdir = get_cr3();
    task->brk = USER_EXEC_ADDR;
    task->exec = USER_EXEC_ADDR;
    task->end = USER_EXEC_ADDR;
    task->text = USER_EXEC_ADDR;
    task->data = USER_EXEC_ADDR;
    task->fpusse = NULL;

    task->umask = 0022;
    task->iexec = NULL;
    task->ipwd = get_root_inode();
    task->iroot = get_root_inode();
    task->iroot->count += 2;
    memset(task->files, 0, sizeof(task->files));

    task->files[STDOUT_FILENO] = &file_table[STDOUT_FILENO];
    task->files[STDIN_FILENO] = &file_table[STDIN_FILENO];
    task->files[STDERR_FILENO] = &file_table[STDERR_FILENO];
    task->files[STDOUT_FILENO]->count++;
    task->files[STDIN_FILENO]->count++;
    task->files[STDERR_FILENO]->count++;

    user_add_task(uid, task);

    task->uid = uid;
    list_append(&schedule_list, &task->rnode);

    for (int i = 0; i < MAXSIG; i++)
    {
        sigaction_t *action = &task->actions[i];
        memset(action, 0, sizeof(sigaction_t));
    }

    task->blocked = 0;

    return task;
}

extern tss_t tss;

void set_gs(u32 gs)
{
    asm volatile("movw %%ax, %%gs\n" :: "a"(gs));
}

void task_activity(task_t *task)
{
    if ((u32)task->pgdir != get_cr3())
    {
        set_cr3((u32)task->pgdir);
    }
    if (task->uid != KERNEL_USER)
    {
        tss.esp0 = (u32)task + PAGE_SIZE;
    }
    if (task->tls)
    {
        tls_update(task);
        set_gs((TLS_IDX << 3) | 0b11);
    }
    if (!task->tls)
    {
        set_gs(USER_DATA_SELECTOR);
    }
}

extern void task_switch(task_t *);

void schedule()
{
    assert(!get_interrupt_state());

    task_t *task = running_task();
    if (task->state == TASK_RUNNING && task->pid != 0)
    {
        list_append(&schedule_list, &task->rnode);
        task->state = TASK_READY;
    }
    task->ticks = task->prio;

    // 任务在内核模式下，可能因为调用某些函数而产生中断（例如ide_udma_write），堵塞了自己
    // 如果这时候又出于某些原因切换了页目录，页目录由 schudule() 保存
    if (get_cr3() != task->pgdir)
        task->pgdir = get_cr3();

    list_node_t *node = list_pop(&schedule_list);

    if (!node && task->state == TASK_READY)
    {
        node = &task->rnode;
        list_remove(node);
        task->state = TASK_RUNNING;
        return;
    }
    if (!node)
    {
        node = &idle_task->rnode;
    }

    task_t *next = element_entry(task_t, rnode, node);
    assert(next);

    assert(next);
    next->state = TASK_RUNNING;

    set_cr0(get_cr0() | CR0_TS);

    task_activity(next);
    task_switch(next);
}

void task_yield()
{
    schedule();
}

int task_block(task_t *task, int state)
{
    task->state = state;

    if (task == running_task())
    {
        schedule();
    }
    else
    {
        list_remove(&task->rnode);
    }

    return task->status;
}

void task_unblock(task_t *task, int status)
{
    task->state = TASK_READY;
    task->status = status;

    // 光速切换
    list_push(&schedule_list, &task->rnode);
}

fd_t task_get_file(task_t *task)
{
    for (size_t i = 0; i < OPEN_NR; i++)
    {
        if (!task->files[i])
            return i;
    }
    panic("%s: no more task files\n", task->name);
}

void task_put_file(task_t *task, fd_t fd)
{
    task->files[fd] = NULL;
}

void task_build_stack(task_t *task)
{
    u32 stack = (u32)task + PAGE_SIZE;

    stack -= sizeof(intr_frame_t);
    intr_frame_t *iframe = (intr_frame_t *)stack;
    iframe->eax = 0;

    stack -= sizeof(task_frame_t);
    task_frame_t *frame = (task_frame_t *)stack;
    frame->eip = interrupt_exit;

    task->stack = (u32)frame;
}

pid_t task_fork()
{
    task_t *task = running_task();
    task_t *child = get_free_task();

    pid_t pid = child->pid;
    memcpy(child, task, PAGE_SIZE);

    child->pid = pid;
    child->ppid = task->pid;

    child->jiffies = 0;
    child->ticks = child->prio;
    child->state = TASK_READY;

    child->pgdir = (u32)copy_pgdir();

    child->vmap = kmalloc(sizeof(bitmap_t));
    child->vmap->bits = (char *)alloc_kpage(1);
    memcpy(child->vmap, task->vmap, sizeof(bitmap_t));
    memcpy(child->vmap->bits, task->vmap->bits, PAGE_SIZE);

    child->ipwd->count++;
    child->iroot->count++;
    if (child->iexec)
        child->iexec->count++;

    for (int i = 0; i < OPEN_NR; i++)
    {
        file_t *file = task->files[i];
        if (!file)
            continue;
        file->count++;
    }

    if (task->fpusse)
    {
        child->fpusse = (u8 *)alloc_kpage(1);
        memcpy(child->fpusse, task->fpusse, PAGE_SIZE);
    }

    if (task->tls)
    {
        child->tls = kmalloc(sizeof(tls_t));
        memcpy(child->tls, child->tls, sizeof(tls_t));
    }

    memset(&child->unode, 0, sizeof(list_node_t));
    memset(&child->rnode, 0, sizeof(list_node_t));
    memset(&child->node, 0, sizeof(list_node_t));

    user_add_task(task->uid, child);
    list_push(&schedule_list, &child->rnode);

    task_build_stack(child);

    return child->pid;
}

void task_exit(int status)
{
    task_t *task = running_task();

    // 禁止 init 进程退出
    assert(task->pid != 1);

    free_kpage((u32)task->vmap->bits, 1);
    kfree(task->vmap);

    free_pgdir();

    iput(task->iroot);
    iput(task->ipwd);
    iput(task->iexec);

    user_remove_task(task);

    task->status = status;

    if (task->fpusse)
        free_kpage((u32)task->fpusse, 1);

    if (task->tls)
        kfree(task->tls);

    for (int i = 0; i < TASK_NR; i++)
    {
        task_t *ptr = task_table[i];
        if (ptr == NULL)
            continue;
        if (ptr->ppid == task->pid)
            ptr->ppid = task->ppid;
    }

    for (int i = 0; i < OPEN_NR; i++)
    {
        file_t *file = task->files[i];
        if (!file)
            continue;
        put_file(file);
    }

    task->state = TASK_DIED;

    if (task->exittid)
    {
        *task->exittid = 0;
        task->exittid = NULL;
    }

    task_t *parent = task_table[task->ppid];
    if (parent->state == TASK_WAITING &&
        (parent->waitpid == -1 || parent->waitpid == task->pid))
    {
        task_unblock(parent, EOK);
    }

    schedule();
}

pid_t task_waitpid(pid_t pid, int *status, int options)
{
    task_t *task = running_task();
    task_t *child = NULL;

    int i = 0;

    if (pid < -1)
        pid = -pid;

    if (pid > 0)
        child = task_table[pid];
    else if (pid != -1)
        return EOF;
    else while (i < TASK_NR)
    {
        task_t *ptr = task_table[i++];
        if (!ptr)
            continue;
        if (ptr->ppid != task->pid)
            continue;
        child = ptr;
        pid = child->pid;
        break;
    }

    if (!child)
        return EOF;
    if (child->state != TASK_READY && options & WUNTRACED)
        return 0;
    if (child->state != TASK_DIED && options & WNOHANG)
        return 0;
    if (child->state != TASK_DIED)
    {
        task->waitpid = pid;
        task_block(task, TASK_WAITING);
    }

    *status = child->status;
    return task_remove(child);
}

int sys_set_thread_area(user_desc_t *desc)
{
    // if (desc->entry_number != -1 && desc->entry_number != TLS_IDX)
    //     return -EINVAL;

    assert(desc->entry_number == -1 || desc->entry_number == TLS_IDX);

    if (desc->base_addr < USER_EXEC_ADDR || desc->base_addr >= USER_MEMORY_END)
        return -EINVAL;

    task_t *task = running_task();
    tls_t *tls = task->tls;

    u32 size = desc->limit + 1;
    if (desc->limit_in_pages)
        size *= PAGE_SIZE;

    DEBUGK("set task area size %d\n", size);

    // 如果没有 TLS 空间，就分配空间
    if (!tls)
    {
        tls = kmalloc(sizeof(tls_t));
        tls->size = size;
        tls->addr = desc->base_addr;
        task->tls = tls;
    }

    // 如果没有制定 entry number 就找一个
    // 但是这里只支持一个 gdt 项
    u32 idx = desc->entry_number;
    if (idx == -1)
        idx = TLS_IDX;
    desc->entry_number = idx;

    tls_update(task);

    intr_frame_t *iframe = (intr_frame_t *)((u32)task + PAGE_SIZE - sizeof(intr_frame_t));
    iframe->gs = (idx << 3) | 0b11;

    return EOK;
}

pid_t sys_set_tid_address(int *ptr)
{
    task_t *task = running_task();
    task->exittid = ptr;
    return task->pid;
}

extern void init_thread();
extern void idle_thread();

void task_init()
{
    memset(task_table, 0, sizeof(task_table));

    list_init(&schedule_list);

    task_t *task = running_task();
    task->state = TASK_INIT;
    task->ticks = 0;
    task->pgdir = get_cr3();

    idle_task = task_create("idle", 1, idle_thread, KERNEL_USER);
    list_remove(&idle_task->rnode);

    task_create("init", 5, init_thread, ROOT_USER);
}